# -*- coding: UTF-8 -*-
"""
    @Author:YTQ
    @Time: 2022/7/20 19:07
    Description: 模型验证代码并模型计算准确率
    
"""

import torch
from torch.autograd import Variable
from torch.utils.data import DataLoader
from dataset import TestDataSet
from CNNet import MyCNN


def test(model_path, image_path, file_path, data_transform, dogAct, catAct, batch_size, workers):
    # load dataset
    # load data
    dataset = TestDataSet(filePath=file_path, image_path=image_path, data_transform=data_transform, dogAct=dogAct, catAct=catAct)
    # 用PyTorch的DataLoader类封装，实现数据集顺序打乱，多线程读取，一次取多个数据等效果
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True,
                            num_workers=workers)

    # load model
    cnn = MyCNN()
    cnn.load_state_dict(torch.load(model_path))

    correct = 0
    total = 0
    with torch.no_grad():  # 进行评测的时候网络不更新梯度
        for images, labels in dataloader:
            images, labels = Variable(images), labels
            outputs = cnn(images)[0]
            predicted = torch.max(outputs.data, 1)[1].data.numpy()
            total += labels.size(0)
            correct += (predicted == labels.numpy()).sum()
        print('Accuracy of the network on the test images: %d %%' % (100 * correct / total))
    # Accuracy of the network on the 100 test images: 90 % 这时我简单训练20轮的模型准确度（仅供参考）
